from state import EditorState
from langchain_core.prompts import ChatPromptTemplate
from model.zhipu import get_model
from utils.get_json import get_json

prompt = ChatPromptTemplate.from_template("""
你是校对员,你需要利用你的专业知识校对提交文档的真实性.
 - 需要校对的文档内容为:{submit}
 - 你需要返回一个json对象,且只有两个key,第一个是pass(是否通过校对考核),第二个是idea(没有通过时,需要给出意见)
""")


def generate_check(state: EditorState):
    title = state['title']
    search = state['search']
    outlines = state['outlines']
    paragraphs = state['paragraphs']
    write_submit = state['write_submit']

    model = get_model()
    chains = prompt | model
    response = chains.invoke({"submit": write_submit})

    print("============ Check Agent ============")
    check_info = get_json(response.content)
    print(check_info)

    return {
        "who": "check",
        "title": title,  # 制定标题
        "search": search,  # 搜集信息
        "outlines": outlines,  # 大纲搜集信息 三级标题
        "paragraphs": paragraphs,  # 段落信息
        "write_submit": write_submit,
        "check_idea": check_info,  # 校对意见
        "comment_idea": {}  # 评论意见
    }


def check_to_write_agent(state: EditorState):
    check_idea = state['check_idea']
    if check_idea['pass']:
        return "comment_agent"
    else:
        return "write_agent"

# if __name__ == '__main__':
#     inputs = {
#         "who": "",
#         "title": "中国的近代历史",  # 制定标题
#         "search": "",  # 搜集信息
#         "outlines": {},  # 大纲搜集信息 三级标题
#         "paragraphs": [],  # 段落信息
#         "check_idea": {},  # 校对意见
#         "comment_idea": {}  # 评论意见
#     }
#
#     outputs = generate_title(inputs)
#     print(outputs)
#     print(outputs['title'])
